import cn.bugstack.chatglm.model.*;
import cn.bugstack.chatglm.session.Configuration;
import cn.bugstack.chatglm.session.OpenAiSessionFactory;
import cn.bugstack.chatglm.session.OpenSession;
import cn.bugstack.chatglm.session.defaults.DefaultOpenAiSessionFactory;
import com.alibaba.fastjson.JSON;
import okhttp3.Response;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.jetbrains.annotations.Nullable;
import org.junit.Before;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.Arrays;
import java.util.concurrent.CountDownLatch;

public class ApiTest {

    private OpenSession openSession;

    private Logger logger = LoggerFactory.getLogger(ApiTest.class);
    @Before
    public void setup() {
        // 创建一个 Configuration 对象
        Configuration configuration = new Configuration();
        // 设置你的 API Secret Key
        configuration.setApiSecretKey("e64bbd89200540b08873d5be17582312.48J1rFfpnm1JIFGc");
        // 设置你的 API Base URL
        configuration.setApiHost("https://open.bigmodel.cn/");
        OpenAiSessionFactory factory = new DefaultOpenAiSessionFactory(configuration);
        openSession = factory.openSession();
    }

    @Test
    public void testChatCompletion() throws IOException {
        ChatCompletionRequest chatCompletionRequest = new ChatCompletionRequest(Model.GLM_4V_PLUS.getModel(),
                Arrays.asList(new ChatCompletionRequest.Message( Role.USER.getRole(),"学习的心得")));
        ChatCompletionResponse chatCompletionResponse = openSession.chatCompletion(chatCompletionRequest);
        logger.info("对话得到的结果:{}",chatCompletionResponse.getChoices().get(0).getMessage().getContent());
    }

    @Test
    public void testCompletion() throws Exception {
        CountDownLatch latch = new CountDownLatch(2);
        ChatCompletionRequest chatCompletionRequest = new ChatCompletionRequest(Model.GLM_4V_PLUS,
                Arrays.asList(new ChatCompletionRequest.Message( Role.USER.getRole(),"学习的心得")));
        EventSource eventSource = openSession.completions(chatCompletionRequest, new EventSourceListener() {
            @Override
            public void onOpen(EventSource eventSource, Response response) {
                logger.info("开始");
                super.onOpen(eventSource, response);
            }

            @Override
            public void onEvent(EventSource eventSource, @Nullable String id, @Nullable String type, String data) {
                CompletionResponse response = JSON.parseObject(data, CompletionResponse.class);
                logger.info("数据:{}",response);
                super.onEvent(eventSource, id, type, data);
            }

            @Override
            public void onClosed(EventSource eventSource) {
                super.onClosed(eventSource);
                logger.info("结束");
                latch.countDown();
            }

            @Override
            public void onFailure(EventSource eventSource, @Nullable Throwable t, @Nullable Response response) {
                super.onFailure(eventSource, t, response);
                logger.info("失败");
                latch.countDown();
            }
        });
        latch.await();
    }

    @Test
    public void testImageCompletion() throws IOException {
        ImageCompletionRequest chatCompletionRequest = new ImageCompletionRequest(Model.COGVIEW_3_PLUS, "一个可爱的泰迪");
        ImageCompletionResponse imageCompletionResponse = openSession.genImages(chatCompletionRequest);
        logger.info("对话得到的结果:{}",imageCompletionResponse.getUrls().get(0));
    }

    @Test
    public void testVideoCompletion() throws Exception {
        VideoCompletionRequest videoCompletionRequest = new VideoCompletionRequest(Model.COGVIDEOX, "老板开心的给员工发工资的时刻");
        VideoCompletionResponse videoCompletionResponse = openSession.genVideos(videoCompletionRequest);
        logger.info("对话得到的结果:{}",videoCompletionResponse.getVideo_result().get(0).getUrl());
    }
}
